'''
二分类决策树模型
'''
from sklearn import tree
import pandas as pd
import numpy as np
from release_code.data_analysis.data_two_metrics import  my_classfication_report, base


# 以医生预测180个实例作为测试集
def read_data():
    test_data = pd.read_csv('D:/lung_cancer/data/divide_csv/two/test.csv')
    train_data = pd.read_csv('D:/lung_cancer/data/divide_csv/two/train.csv')
    test_features = []
    train_features = []
    test_labels = []
    train_labels = []
    for i in range(len(test_data)):
        one_feature = [test_data['z'][i], test_data['x'][i], test_data['y'][i], test_data['r'][i],
                       test_data['patientWeight'][i], test_data['patientSex'][i], test_data['patientAge'][i],
                       test_data['patientSize'][i], test_data['local_suvmax'][i], test_data['local_suvmin'][i],
                       test_data['local_suvavg'][i], test_data['local_suvstd'][i], test_data['local_suvvar'][i]]

        test_features.append(one_feature)
        test_labels.append(test_data['cancer_type'][i] - 1)

    for j in range(len(train_data)):
        one_feature = [train_data['z'][j], train_data['x'][j], train_data['y'][j], train_data['r'][j],
                       train_data['patientWeight'][j], train_data['patientSex'][j], train_data['patientAge'][j],
                       train_data['patientSize'][j], train_data['local_suvmax'][j], train_data['local_suvmin'][j],
                       train_data['local_suvavg'][j],  train_data['local_suvstd'][j], train_data['local_suvvar'][j]]

        train_features.append(one_feature)
        train_labels.append(train_data['cancer_type'][j] - 1)

    X_train = np.asarray(train_features, dtype=np.float)
    X_test = np.asarray(test_features, dtype=np.float)
    y_train = np.asarray(train_labels, dtype=np.int)
    y_test = np.asarray(test_labels, dtype=np.int)

    return X_train, X_test, y_train, y_test


def train():
    train_features, test_features, train_labels, test_labels = read_data()

    clf = tree.DecisionTreeClassifier(criterion='entropy',
                                      max_features=None,
                                      max_depth=5,
                                      min_samples_split=2,
                                      class_weight='balanced')
    clf.fit(train_features, train_labels)

    preds = clf.predict(test_features)

    tp, tn, fp, fn = base(test_labels, preds, 0.5)
    # print(tp, tn, fp, fn)
    tpr = float(tp) / (tp + fn)
    fpr = float(fp) / (fp + tn)
    print('二分类决策树ROC结果横坐标（fpr）：%f' % fpr)
    print('二分类决策树ROC结果纵坐标（tpr）：%f' % tpr)

    precision = float(tp) / (tp + fp)
    recall = float(tp) / (tp + fn)
    print('二分类决策树P-R结果横坐标（recall）：%f' % recall)
    print('二分类决策树P-R结果纵坐标（precision）：%f' % precision)


if __name__ == '__main__':
    train()
